神经网络量化(TensorRT 和 IAO)

量化简单来说就是使用一种低精度的方式来作为存储和计算的数值表示方式。

常见方法:NVIDIA 的 TensorRT 和 Google 的 IAO 量化方法

TensorRT

Tensor Values = FP32 scale factor * INT8 array + FP32 bias

NVIDIA 研究人员通过实验发现,其实我们并不需要在量化的时候加上偏置

一种理解:主要是因为偏置对于一组数值而言,其改变的是数值的分布位置,但是当前神经网络的归一化操作很多,因此可以去掉偏置。

那么,对于 A=scaleAQA+biasAB=scaleBQB+biasB
A×B=scaleA×scaleB×QA×QB

Pasted image 20230530112001.png

一些网络中激活值的分布统计:

Pasted image 20230530112140.png

KLD

信息熵:H(X,p)=xXp(x)logp(x)

交叉熵:H(X,p,q)=xXp(x)logq(x)

相对熵:D(p||q)=H(X,p,q)H(X,p)=xXp(x)(logq(x)logp(x))

TensorRT(NCNN) 中寻找阈值和计算 scale factor 的流程:

  1. 准备一个校准数据集,覆盖模型的使用场景即可,数量不需要很多。
  2. 将数据集的每张图片都通过模型做一次预测,在这个过程中,对所有要量化的层的 FP32 数值分布进行统计,得到 |max|。
    • 量化是针对每个 channel 单独做的,因此卷积层的每个 channel 都是单独统计、计算和量化的,后续操作也是。
  3. 将 0~|max| 分成 n 个 bin,然后再次遍历所有图片,让每个量化层中的数值落到其属于的 bin 中,统计每个 bin 的数目。
    • |max| / n 就可以得到每个 bin 的宽度 w,因此 就分为了 0~w, w~2*w…(n-1)*w~|max| 总共 n 个 bin。
    • 对每个数值按照其绝对值分到不同 bin 中。
    • TensorRT 官方使用的 n 是 2048,mxnet 是 4096,n 越大越好,但是计算量会上升。
  4. 遍历第 128~n 个 bin:
    • 以当前 bin 的中值作为阈值 T,在其上做截断,将大于当前 bin 以外的 bin 的统计数目加到当前 bin 上,这一步是为了减少直接抹去这些数值带来的影响;
    • 计算此时的概率分布 P,每个 bin 的概率就是其统计数目除以数值的总数;
    • 创建一个新的编码 Q,其长度是128,其元素的值就是 P 量化后的 INT8 数值(正数是0~+127,负数是-128~-1);
    • 因为 Q 分布只有 128 个编码,为了计算交叉熵,将其扩展到和 P 同样的长度;
      Pasted image 20230611104855.png
    • 计算 P 和 Q 的相对熵,记录当前相对熵和阈值 T
  5. 选择最小的相对熵和其对应阈值 T,计算 scale factor = T / 127
    • 实际代码中使用的是 scale factor = 127 / T,这样 FP32 到 INT8 量化的时候可以使用乘法而不是除法。
  6. 对每个 bin,取其中值作为这个 bin 当前的 FP32 表示,然后除以 scale factor,然后四舍五入,就得到了其量化后的 INT8 数值,将这个 bin 中所有的 FP32 数值都映射为这个 INT8 表示的数。多个 bin 可能映射为同一个 INT8 数字。

NCNN 量化实现

https://github.com/Tencent/ncnn/tree/master/tools/quantize

NCNN 对 conv 和 fc 进行量化 (input_data 和 weights 分开计算 scale factor)

Pasted image 20230530172318.png
Pasted image 20230530172324.png
首先通过 校准数据集 生成量化表(scale factor,离线生成, ncnn2table.cpp)

量化 weights:只需要使用非饱和量化(scale factor = |max|),离线完成

量化 input_data:用校准数据集定下最佳阈值(这个过程比较花费时间,所以提前预处理),保存在量化表中;数据不确定——只能在线完成

requantize

当一个Conv1后面紧跟着另一个Conv2时,NCNN会进行 requantize 的操作:在得到Conv1的INT32输出后,会顺手帮Conv2做量化,得到Conv2的INT8输入,节省一次内存读写。

requantize 相当于替换了原来的 dequantize+add_bias+下一层的quantize

Pasted image 20230530172453.png

目前NCNN只有两种情况会用到requantize:

(1) Conv1 -> ReLU -> Conv2,且Conv1和Conv2都用int8
(2) Conv1 -> ReLU -> Split -> Conv2_1
-> Conv2_2
...
-> Conv2_x,且这些Conv都用int8

IAO

量化公式:Q=RS+Z
其中,R 表示真实的浮点值,Q 表示量化后的定点值,Z 表示 0 浮点值对应的量化定点值,也称零点漂移,S 则为量化的 scale factor,S 和 Z 的求值公式如下:(饱和映射)

S=RmaxRminQmaxQminZ=QmaxRmaxS

这里的 S 和 Z 均是量化参数,S 是 FP32 类型,而 Z 是 INT8 类型。 Q 和 R 均可由公式进行求值,不管是量化后的 Q 还是反推求得的浮点值 R,如果它们超出各自可表示的最大范围,那么均需要进行截断处理。

IAO对应的矩阵乘法公式:

S3(q3(i,k)Z3)=j=1NS1(q1(i,j)Z1)S2(q2(j,k)Z2)q3(i,k)=Z3+Mj=1N(q1(i,j)Z1)(q2(j,k)Z2)

其中 M=S1S2S3 ,整个式子中只有 M 是浮点数,且实际中 M(0,1) ,所以考虑将 M 也用整数表示:M0=2nM ,其中 M0 in(0.5,1]

IAO 量化训练

IAO 方法将权重、激活值及输入值均全部做 INT8 量化,并且将所有模型运算操作置于 INT8 下进行执行,不可避免地会产生误差,所以需要实施量化精度补偿训练。

  1. 输入 量化的特征图 lhs_quantized_val,INT8 类型,偏移量 lhs_zero_point,INT32 类型;
  2. 输入 量化的卷积核 rhs_quantized_val,INT8 类型,偏移量 rhs_zero_point,INT32 类型;
  3. 转换 INT8 到 INT32类型;
  4. 每一块卷积求和(INT32 乘法求和有溢出风险,可换成固定点小数乘法);
    • int32_accumulator += (lhs_quantized_val(i, j) - lhs_zero_point) * (rhs_quantized_val(j, k) - rhs_zero_point)
  5. 输入 量化的乘子 quantized_multiplier(M0,INT32 类型小数)和右移次数记录 right_shift(n),INT 类型;
    • real_multiplier = (S_in * S_W / S_out) * (2^n)
    • quantized_multiplier = round(real_multiplier * (1 << 31))
  6. 计算乘法,得到 INT32 类型的结果 (INT32乘法有溢出风险,可换成固定点小数乘法);
    • int32_result = quantized_multiplier * int32_accumulator
  7. 再左移动 right_shift 位还原,得到 INT32 的结果;
  8. 最后再加上结果的偏移量 result_zero_point;
  9. 将 INT32 类型结果强制转换到 INT8 类型,就得到了量化后的矩阵计算结果;
  10. 之后再反量化到浮点数,更新统计输出值分布信息 Rmax,,Rmin
  11. 再量化回 INT8;
  12. 之后经过量化激活层;
  13. 最后反量化到浮点数,本层网络输出;
  14. 进入下一层,循环执行 1~12 步骤

Pasted image 20230611112339.png

refer